# intelligent_sop_cms/decision_engine.py
from datetime import datetime
from typing import List, Dict, Any, Optional
from enum import Enum
import json
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
import joblib


class DecisionType(Enum):
    SOP_SELECTION = "sop_selection"
    RISK_ASSESSMENT = "risk_assessment"
    RESOURCE_ALLOCATION = "resource_allocation"
    WORKFLOW_OPTIMIZATION = "workflow_optimization"
    EMERGENCY_RESPONSE = "emergency_response"


class AIDecisionEngine:
    """自主决策引擎"""

    def __init__(self):
        self.ml_model = None
        self.decision_rules = {}
        self.historical_data = []
        self.learning_enabled = True
        self.confidence_threshold = 0.7

    def initialize(self):
        """初始化决策引擎"""
        self._load_decision_rules()
        self._train_initial_model()
        print("AI决策引擎初始化完成")

    def make_autonomous_decision(self, context: Dict[str, Any], decision_type: DecisionType) -> Dict[str, Any]:
        """做出自主决策"""
        # 分析上下文
        context_analysis = self._analyze_context(context)

        # 风险评估
        risk_level = self._assess_risk(context, decision_type)

        # 基于规则和机器学习做出决策
        if self.learning_enabled and self.ml_model:
            ml_decision = self._ml_based_decision(context, decision_type)
            rule_decision = self._rule_based_decision(context, decision_type)

            # 融合决策
            final_decision = self._fuse_decisions(ml_decision, rule_decision, context_analysis)
        else:
            final_decision = self._rule_based_decision(context, decision_type)

        # 记录决策
        decision_record = {
            "timestamp": datetime.now(),
            "decision_type": decision_type.value,
            "context": context,
            "decision": final_decision,
            "risk_level": risk_level,
            "confidence": final_decision.get("confidence", 0.0)
        }
        self.historical_data.append(decision_record)

        return final_decision

    def _analyze_context(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """分析决策上下文"""
        analysis = {
            "complexity": self._calculate_complexity(context),
            "urgency": context.get("urgency", "normal"),
            "available_resources": len(context.get("resources", [])),
            "constraints": len(context.get("constraints", [])),
            "historical_similarity": self._find_similar_historical_cases(context)
        }
        return analysis

    def _assess_risk(self, context: Dict[str, Any], decision_type: DecisionType) -> str:
        """风险评估"""
        risk_factors = {
            "sop_complexity": len(context.get("required_steps", [])),
            "human_factors": context.get("operator_experience", "beginner"),
            "equipment_risk": context.get("equipment_condition", "unknown"),
            "environmental_factors": context.get("environment", "normal")
        }

        risk_score = 0
        if risk_factors["sop_complexity"] > 10:
            risk_score += 2
        if risk_factors["human_factors"] == "beginner":
            risk_score += 2
        if risk_factors["equipment_risk"] == "poor":
            risk_score += 3
        if risk_factors["environmental_factors"] == "hazardous":
            risk_score += 3

        if risk_score >= 6:
            return "high"
        elif risk_score >= 3:
            return "medium"
        else:
            return "low"

    def _ml_based_decision(self, context: Dict[str, Any], decision_type: DecisionType) -> Dict[str, Any]:
        """基于机器学习的决策"""
        try:
            # 特征工程
            features = self._extract_features(context, decision_type)

            # 预测
            prediction = self.ml_model.predict([features])[0]
            probability = np.max(self.ml_model.predict_proba([features])[0])

            return {
                "decision": prediction,
                "confidence": probability,
                "method": "machine_learning",
                "features_used": len(features)
            }
        except Exception as e:
            print(f"ML决策失败: {e}")
            return {"decision": "fallback", "confidence": 0.0, "method": "fallback"}

    def _rule_based_decision(self, context: Dict[str, Any], decision_type: DecisionType) -> Dict[str, Any]:
        """基于规则的决策"""
        decision_rules = self.decision_rules.get(decision_type.value, [])

        for rule in decision_rules:
            if self._evaluate_rule(rule["condition"], context):
                return {
                    "decision": rule["action"],
                    "confidence": rule.get("confidence", 0.8),
                    "method": "rule_based",
                    "rule_id": rule["id"]
                }

        # 默认决策
        return {
            "decision": "proceed_with_caution",
            "confidence": 0.5,
            "method": "default"
        }

    def _fuse_decisions(self, ml_decision: Dict, rule_decision: Dict, context_analysis: Dict) -> Dict[str, Any]:
        """融合多个决策结果"""
        ml_confidence = ml_decision.get("confidence", 0.0)
        rule_confidence = rule_decision.get("confidence", 0.0)

        # 根据复杂性调整权重
        complexity = context_analysis["complexity"]
        if complexity > 0.7:  # 高复杂性，更依赖ML
            ml_weight = 0.7
            rule_weight = 0.3
        else:  # 低复杂性，更依赖规则
            ml_weight = 0.3
            rule_weight = 0.7

        final_confidence = (ml_confidence * ml_weight + rule_confidence * rule_weight)

        # 选择置信度更高的决策
        if ml_confidence >= rule_confidence:
            final_decision = ml_decision["decision"]
            method = "machine_learning"
        else:
            final_decision = rule_decision["decision"]
            method = "rule_based"

        return {
            "decision": final_decision,
            "confidence": final_confidence,
            "method": f"fused_{method}",
            "components": {
                "ml_decision": ml_decision,
                "rule_decision": rule_decision
            }
        }

    def _load_decision_rules(self):
        """加载决策规则"""
        self.decision_rules = {
            "sop_selection": [
                {
                    "id": "rule_001",
                    "condition": "context.urgency == 'high' and context.risk == 'low'",
                    "action": "select_standard_sop",
                    "confidence": 0.9
                },
                {
                    "id": "rule_002",
                    "condition": "context.operator_experience == 'beginner'",
                    "action": "select_detailed_sop_with_guidance",
                    "confidence": 0.8
                }
            ],
            "risk_assessment": [
                {
                    "id": "rule_101",
                    "condition": "context.equipment_condition == 'poor'",
                    "action": "require_safety_review",
                    "confidence": 0.95
                }
            ]
        }

    def _train_initial_model(self):
        """训练初始机器学习模型"""
        # 这里使用模拟数据，实际中应该使用历史决策数据
        X = np.random.rand(100, 10)  # 10个特征
        y = np.random.randint(0, 3, 100)  # 3种决策

        self.ml_model = RandomForestClassifier(n_estimators=100)
        self.ml_model.fit(X, y)

        # 保存模型
        joblib.dump(self.ml_model, 'decision_model.pkl')

    def _extract_features(self, context: Dict, decision_type: DecisionType) -> List[float]:
        """从上下文中提取特征"""
        # 简化的特征提取，实际中应该更复杂
        features = [
            len(context.get("resources", [])),
            len(context.get("constraints", [])),
            1 if context.get("urgency") == "high" else 0,
            len(context.get("historical_cases", [])),
            context.get("time_constraint", 0) / 3600,  # 小时
        ]

        # 填充到10个特征
        while len(features) < 10:
            features.append(0.0)

        return features[:10]

    def _calculate_complexity(self, context: Dict) -> float:
        """计算任务复杂性"""
        complexity_score = 0.0

        # 基于步骤数量
        steps = len(context.get("required_steps", []))
        complexity_score += min(steps / 20.0, 1.0) * 0.4

        # 基于资源约束
        resources = len(context.get("resources", []))
        constraints = len(context.get("constraints", []))
        complexity_score += min(constraints / max(resources, 1), 1.0) * 0.3

        # 基于时间压力
        if context.get("urgency") == "high":
            complexity_score += 0.3

        return min(complexity_score, 1.0)

    def _evaluate_rule(self, condition: str, context: Dict) -> bool:
        """评估规则条件"""
        # 简化的规则评估，实际中应该使用规则引擎
        try:
            # 这里使用简单的字符串匹配，实际应该用更复杂的逻辑
            if "urgency == 'high'" in condition and context.get("urgency") == "high":
                return True
            if "operator_experience == 'beginner'" in condition and context.get("operator_experience") == "beginner":
                return True
            return False
        except:
            return False

    def _find_similar_historical_cases(self, context: Dict) -> List[Dict]:
        """查找相似的历史案例"""
        similar_cases = []
        for case in self.historical_data[-100:]:  # 最近100个案例
            similarity = self._calculate_case_similarity(context, case["context"])
            if similarity > 0.6:
                similar_cases.append({
                    "case": case,
                    "similarity": similarity
                })

        return sorted(similar_cases, key=lambda x: x["similarity"], reverse=True)[:5]

    def _calculate_case_similarity(self, context1: Dict, context2: Dict) -> float:
        """计算案例相似度"""
        common_keys = set(context1.keys()) & set(context2.keys())
        if not common_keys:
            return 0.0

        similarity_score = 0.0
        for key in common_keys:
            if context1[key] == context2[key]:
                similarity_score += 1.0

        return similarity_score / len(common_keys)